#include "../common/mmo.h"
#include "../common/showmsg.h"
#include "../common/malloc.h"
#include "../common/timer.h"
#include "../common/socket.h"

#include <stdio.h>
#include <stdlib.h>
#include <signal.h>
#include <string.h>
#ifndef _WIN32
#include <unistd.h>
#endif
#include <openssl/aes.h>

struct harmony_data {
  struct {
    int crypted; // 0 - ??, 1 - yes, 2 - no
    struct s_harmony_packet_crypt {
      uint8 *key;
      size_t pos;
    } in, out;
    uint8 *challenge;
    size_t challenge_size;
  } packet;
  struct {
    bool enabled;
    uint32 val;
    size_t pos;
  } idcrypt;
};

uint16 out_translate_table[65536];

uint8* idcrypt_key;
size_t idcrypt_key_size;
uint32 idcrypt_default_val;

uint8* crypt_key;
size_t crypt_key_size;

uint8 login_crypt_key[32];

// ---

void read_databin(void);

// ----

void harmony_core_init() {
  if( sizeof(out_translate_table) != 2*65536 ) {
    ShowError("Invalid harmony table size.\n");
    exit(EXIT_FAILURE);
  }

  // Read db
  read_databin();
}

void harmony_core_final() {
  aFree(crypt_key);
  aFree(idcrypt_key);
}

void harmony_process_new(int fd) {
  CREATE(session[fd]->harm, struct harmony_data, 1);
}

void harmony_process_del(int fd) {
  if( session[fd]->harm->packet.challenge )
    aFree(session[fd]->harm->packet.challenge);
  if( session[fd]->harm->packet.crypted == 1 ) {
    aFree(session[fd]->harm->packet.in.key);
    aFree(session[fd]->harm->packet.out.key);
  }
  aFree(session[fd]->harm);
}

void harmony_crypt(int fd, struct s_harmony_packet_crypt *c, uint8 *buf, size_t len) {
  size_t i;
  bool outgoing = (c == &session[fd]->harm->packet.out);
  
  for(i = 0; i < len; i++) {
    uint8 plain;
    
    if( outgoing )
      plain = buf[i];
    buf[i] ^= c->key[c->pos];
    if( !outgoing )
      plain = buf[i];
    c->pos = (c->pos + 1) % crypt_key_size;
    
    if( c->pos % 79 == 4 ) {
      size_t j;
      
      for(j = 8; j > 0; j--) {
        AES_KEY aeskey;
        size_t k;
        size_t k_pos = (c->pos + plain * 189 + c->key[c->pos] * 13) % (crypt_key_size - j * 16);
        size_t p_pos = (c->pos + c->key[c->pos] * plain + 4) % (crypt_key_size - 16);
        
        if( abs((int)k_pos - (int)p_pos) < 32 )
          break;
        
        AES_set_encrypt_key(&c->key[k_pos], 32*8, &aeskey);
        AES_ecb_encrypt(&c->key[p_pos], &c->key[p_pos], &aeskey, AES_ENCRYPT);
        
        for(k = 0; k < 16; k++)
          c->key[p_pos + k] ^= c->key[k_pos + k];
        break;
      }
    }
  }
}

void harmony_gen_key(uint8 *key, uint32 init_val) {
  uint32 poly, v;
  uint32 *tab = (uint32*)key;
  size_t i, j;
  
  poly = 0xa9384445 ^ init_val;
  for(i = 0; i < (crypt_key_size / sizeof(uint32)); i++) {
    v = i;
    for(j = 8; j > 0; j--) {
      if( v & 1 )
        v = (v >> 1) ^ poly;
      else
        v >>= 1;
    }
    tab[i] = v ^ 0xd71a89b0;
  }
  
  for(i = 0; i < crypt_key_size; i++) {
    if( i % 2 )
      key[i] ^= crypt_key[i];
    else
      key[i] += crypt_key[i];
  }
  
}

void harmony_process_recv(int fd, uint8 *buf, size_t len) {
  struct harmony_data *sd = session[fd]->harm;
  
  if( sd->packet.crypted == 0 ) {
    uint8 *challenge;
    
    if( !sd->packet.challenge ) {
      if( RFIFOREST(fd) < 2 )
        return; // need packet id
      if( RFIFOW(fd, 0) != 0x0072 ) {
        sd->packet.crypted = 2;
        return;
      }
    }
    
    if( sd->packet.challenge ) {
      size_t l = min(39-sd->packet.challenge_size, RFIFOREST(fd));
      memcpy(sd->packet.challenge + sd->packet.challenge_size, RFIFOP(fd, 0), l);
      sd->packet.challenge_size += l;
      buf += l;
      len -= l;
      RFIFOSKIP(fd, l);
      if( sd->packet.challenge_size < 39 )
        return; // not done
      challenge = sd->packet.challenge;
    } else if( !sd->packet.challenge && RFIFOREST(fd) >= 39 ) {
      challenge = RFIFOP(fd, 0);
      buf += 39;
      len -= 39;
      RFIFOSKIP(fd, 39);
    } else {
      CREATE(sd->packet.challenge, uint8, 39);
      memcpy(sd->packet.challenge, RFIFOP(fd, 0), RFIFOREST(fd));
      sd->packet.challenge_size += RFIFOREST(fd);
      RFIFOSKIP(fd, RFIFOREST(fd));
      return; // not done
    }
    sd->packet.crypted = 1;
    
    CREATE(sd->packet.in.key, uint8, crypt_key_size);
    CREATE(sd->packet.out.key, uint8, crypt_key_size);
    sd->packet.in.pos = 0;
    sd->packet.out.pos = 0;
    
    harmony_gen_key(sd->packet.in.key, WBUFL(challenge, 3) ^ WBUFL(challenge, 8));
    memcpy(sd->packet.out.key, sd->packet.in.key, crypt_key_size);
    
    aFree(sd->packet.challenge);
    sd->packet.challenge = NULL;
    sd->packet.challenge_size = 0;
  }
  
  if( sd->packet.crypted == 1 ){
    harmony_crypt(fd, &sd->packet.in, buf, len);
  }
}

void harmony_process_send(int fd, uint8 *buf, size_t len) {
  struct harmony_data *sd = session[fd]->harm;
  
  // WBUFW(buf, 0) = out_translate_table[WBUFW(buf, 0)];
  
  if( sd->packet.crypted == 1 ) {
    harmony_crypt(fd, &sd->packet.out, buf, len);
  }
}

bool harmony_packetcrypt_is_enabled(int fd) {
  return (session[fd]->harm->packet.crypted == 1);
}

void harmony_packetcrypt_enable(int fd) {
  session[fd]->harm->packet.crypted = 1;
}

void harmony_idcrypt_enable(int fd) {
  session[fd]->harm->idcrypt.val = idcrypt_default_val;
  session[fd]->harm->idcrypt.pos = 0;
  session[fd]->harm->idcrypt.enabled = true;
}

uint16 harmony_idcrypt(int fd, int apply) {
  struct harmony_data *sd = session[fd]->harm;
  uint32 v;
  uint16 id;
  size_t pos;
  
  if( !sd->idcrypt.enabled )
    return RFIFOW(fd, 0);
  
  v = sd->idcrypt.val * 0xDCB7BD47 + idcrypt_key[sd->idcrypt.pos] * 3;
  id = RFIFOW(fd, 0) ^ (v & 0xFFFF);
  pos = (sd->idcrypt.pos + 1 + (id >> 1)) % idcrypt_key_size;
  
  if( apply ) {
    sd->idcrypt.val = v;
    sd->idcrypt.pos = pos;
    RFIFOW(fd, 0) = id;
  }
  
  return id;
}

void read_databin(void) {
  FILE *fp;
  int need = (1 | 2 | 4 | 8 | 16);
  
  fp = fopen("./harmony/data.bin", "rb");
  if( !fp ) {
    ShowError("Failed to read harmony/data.bin\n");
    exit(EXIT_FAILURE);
  }
  while( !feof(fp) ) {
    uint8 op = fgetc(fp);
    
    if( feof(fp) ) break;
    
    if( op > 30 || (need & (1 << (op - 1))) == 0 )
      goto __invalid;
    need ^= (1 << (op - 1));
    
    switch( op ) {
    case 1:
      idcrypt_key_size = 0; // x64..
      if( !fread((void*)&idcrypt_key_size, 4, 1, fp) )
        goto __invalid;
      CREATE(idcrypt_key, uint8, idcrypt_key_size);
      if( !fread(idcrypt_key, idcrypt_key_size, 1, fp) )
        goto __invalid;
      break;
    case 2:
      if( !fread((void*)&idcrypt_default_val, 4, 1, fp) )
        goto __invalid;
      break;
    case 3:
      crypt_key_size = 0; // x64..
      if( !fread((void*)&crypt_key_size, 4, 1, fp) )
        goto __invalid;
      CREATE(crypt_key, uint8, crypt_key_size);
      if( !fread(crypt_key, crypt_key_size, 1, fp) )
        goto __invalid;
      break;
    case 4:
      if( !fread((void*)out_translate_table, sizeof(out_translate_table), 1, fp) )
        goto __invalid;
      break;
	case 5:
	  if( !fread((void*)login_crypt_key, 32, 1, fp) )
		goto __invalid;
    }
  }
  
  if( need != 0 )
    goto __invalid;
  
  fclose(fp);
  return;
  
__invalid:
  ShowError("Invalid harmony/data.bin.\n");
  exit(EXIT_FAILURE);
}

